import numpy
import requests
# from transformers import sentence_transformers 考虑到服务器性能我直接调用的api


# embedding 计算
def siliconflow_embedding(texts: list,model_name: str="netease-youdao/bce-embedding-base_v1"):
    """
        用于计算文本的embedding
        texts: list, 输入文本列表
        model_name: str, 模型名称 netease-youdao/bce-embedding-base_v1、BAAI/bge-m3
        output: list, 返回文本的embedding
    """
    url = "https://api.siliconflow.cn/v1/embeddings"
    payload = {
        "model": model_name,
        "input": texts,
        "encoding_format": "float"
    }
    headers = {
        "Authorization": "Bearer sk-nihhjskdaommndqhohdrneftvrfukgavarvteuearcrdjgqw",
        "Content-Type": "application/json"
    }
    response = requests.request("POST", url, json=payload, headers=headers)
    # print(response.json())
    embeddings = response.json()["data"]
    embeddings = [item["embedding"] for item in embeddings]
    return embeddings



if __name__ == '__main__':
    texts = ["Hello, my name is John.", "I am a student."]
    embeddings = siliconflow_embedding(texts)
    print(numpy.array(embeddings).shape)